# 将PDF文件存入Milvus并使用DeepSeek模型进行问答的功能。该脚本整合了前两个脚本的核心功能，并添加了PDF处理和问答系统
# PDF处理与向量化 1、使用PyPDFLoader加载PDF文档 2、采用RecursiveCharacterTextSplitter分割文本（块大小1000，重叠200） 3、记录每个文本片段的来源和页码信息
# 向量数据库管理 1、自动检测并重建Milvus集合 2、优化集合结构（包含页码元数据） 3、使用DeepSeek-R1:1.5b生成1536维向量 4、创建IVF_FLAT索引加速检索
# 问答系统实现
#   双模型架构：嵌入模型：处理文本向量化  生成模型：生成自然语言回答
#   检索增强生成（RAG）模式：先检索最相关的5个文本片段   将检索结果作为上下文生成回答
#   显示回答来源（页码和文本片段）
# 交互式界面 1、实时问答反馈 2、显示响应时间和来源信息 3、支持自然退出命令


from config import envConfig, config
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings
from pymilvus import (
    connections, utility, Collection,
    FieldSchema, CollectionSchema, DataType
)
import logging
import time
from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import Milvus
import os
import json

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('insurance_qa_system')

# 全局配置
PDF_PATH = r"D:\PyCharmMiscProject\pythonProjectYA\中间结果\中间结果2\中荷传世赢家（臻享版）终身寿险产品说明.pdf"
COLLECTION_NAME = "insurance_manual"
EMBEDDING_MODEL = "deepseek-r1:1.5b"
LLM_MODEL = "deepseek-r1:1.5b"


def create_collection(collection_name, dim=1536):
    """创建Milvus集合（使用DeepSeek嵌入模型）"""
    # 简化字段定义 - 移除多余的字段
    fields = [
        FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(name="metadata", dtype=DataType.JSON)  # 将page合并到metadata中
    ]

    # 创建schema
    schema = CollectionSchema(fields, description="保险产品说明文档向量库")

    # 创建集合
    collection = Collection(
        name=collection_name,
        schema=schema,
        using="default"
    )

    # 创建索引
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128}  # 减少nlist以加快索引
    }
    collection.create_index("vector", index_params)
    logger.info(f"集合 {collection_name} 创建成功，维度: {dim}")
    return collection


def load_and_split_pdf(pdf_path):
    """加载PDF文件并分割文本"""
    logger.info(f"开始加载PDF文件: {pdf_path}")

    # 使用PyPDFLoader加载文档
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()

    # 设置文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,  # 减少块大小
        chunk_overlap=100,  # 减少重叠
        length_function=len,
        is_separator_regex=False,
    )

    # 分割文档
    chunks = text_splitter.split_documents(documents)
    logger.info(f"PDF分割完成，共 {len(chunks)} 个文本片段")

    # 确保每个chunk都有page元数据
    for chunk in chunks:
        if 'page' not in chunk.metadata:
            chunk.metadata['page'] = 0
        # 添加文件名到元数据
        chunk.metadata['source'] = os.path.basename(pdf_path)

    return chunks


def vectorize_pdf(pdf_path, collection_name):
    """将PDF内容向量化并存储到Milvus"""
    # 1. 连接Milvus
    milvus_conn = {
        "host": envConfig.MILVUS_LOCAL_URI.split("://")[-1].split(":")[0],
        "port": "19530",
        "user": envConfig.MILVUS_USER,
        "password": envConfig.MILVUS_PASSWORD
    }

    try:
        connections.connect(
            alias="default",
            host=milvus_conn["host"],
            port=milvus_conn["port"],
            user=milvus_conn["user"],
            password=milvus_conn["password"],
            secure=False
        )
        logger.info(f"成功连接Milvus: {milvus_conn['host']}:{milvus_conn['port']}")
    except Exception as e:
        logger.error(f"Milvus连接失败: {str(e)}")
        return None

    # 2. 检查并重建集合
    if utility.has_collection(collection_name):
        logger.warning(f"集合 {collection_name} 已存在，正在删除重建...")
        utility.drop_collection(collection_name)
        time.sleep(1)

    # 初始化嵌入模型
    embeddings = OllamaEmbeddings(
        model=EMBEDDING_MODEL,
        base_url=envConfig.OLLAMA_URL
    )

    # 测试嵌入模型获取维度
    try:
        test_embed = embeddings.embed_query("测试文本")
        dim = len(test_embed)
        logger.info(f"嵌入模型维度: {dim}")
    except Exception as e:
        logger.error(f"嵌入模型测试失败: {e}")
        return None

    # 3. 创建新集合
    collection = create_collection(collection_name, dim)

    # 4. 加载和分割PDF
    chunks = load_and_split_pdf(pdf_path)

    # 5. 向量化并存储
    logger.info("开始向量化并存储文档...")
    start_time = time.time()

    # 使用LangChain的Milvus集成
    try:
        vector_db = Milvus.from_documents(
            documents=chunks,
            embedding=embeddings,
            collection_name=collection_name,
            connection_args=milvus_conn,
            vector_field="vector",
            text_field="text",
            index_params={"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 128}},
            drop_old=True  # 确保删除旧集合
        )
        logger.info(f"文档存储完成! 耗时: {time.time() - start_time:.2f}秒")
        logger.info(f"共存储 {len(chunks)} 个文本片段")
        return vector_db
    except Exception as e:
        logger.error(f"向量化过程中出错: {str(e)}")
        return None


def initialize_qa_system(vector_db):
    """初始化问答系统"""
    # 1. 初始化语言模型
    llm = Ollama(
        model=LLM_MODEL,
        base_url=envConfig.OLLAMA_URL,
        temperature=0.2,
        num_ctx=4096,
        system="你是一个专业的保险顾问，根据提供的保险产品说明文档回答问题。"
    )

    # 2. 创建检索式QA链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_db.as_retriever(search_kwargs={"k": 2}),  # 进一步减少检索结果数量
        return_source_documents=True,
        verbose=False
    )

    logger.info("问答系统初始化完成")
    return qa_chain


def interactive_qa(qa_chain):
    """交互式问答循环"""
    print("\n===== 保险产品问答系统 =====")
    print("输入关于保险产品的问题，输入 'exit' 退出")
    print("注意: 系统基于《中荷传世赢家（臻享版）终身寿险产品说明》文档回答")

    while True:
        question = input("\n问题: ").strip()
        if not question:
            continue
        if question.lower() in ["exit", "quit"]:
            break

        start_time = time.time()
        try:
            # 使用更简洁的问题格式
            result = qa_chain.invoke({"query": question})

            # 显示回答
            print("\n回答:", result["result"])

            # 显示来源
            print("\n来源:")
            for i, doc in enumerate(result["source_documents"]):
                page_num = doc.metadata.get('page', 0) + 1
                source = doc.metadata.get('source', '未知来源')
                print(f"{i + 1}. 页码 {page_num}: {doc.page_content[:100]}...")

            print(f"\n响应时间: {time.time() - start_time:.2f}秒")

        except Exception as e:
            logger.error(f"问答过程中出错: {str(e)}")
            print("抱歉，处理问题时出错，请重试")


if __name__ == "__main__":
    # 确保PDF文件存在
    if not os.path.exists(PDF_PATH):
        logger.error(f"PDF文件不存在: {PDF_PATH}")
        exit(1)

    print(f"开始处理保险文档: {os.path.basename(PDF_PATH)}")

    # 步骤1: 向量化PDF文档
    vector_db = vectorize_pdf(PDF_PATH, COLLECTION_NAME)

    if vector_db:
        # 步骤2: 初始化问答系统
        qa_chain = initialize_qa_system(vector_db)

        # 步骤3: 启动交互式问答
        interactive_qa(qa_chain)

    print("\n程序结束。。")